import asyncio
import base64
from typing import Dict, Any, Optional
import logging
from datetime import datetime, timedelta, timezone
import aiohttp
import json

from .base import BaseExecutor
from ..entities import (
    CodeExecutionRequest, 
    ExecutionResult, 
    TaskStatus, 
    CodeType,
    ExecutorConfig
)

logger = logging.getLogger(__name__)

class ContainerExecutor(BaseExecutor):
    """基于Kubernetes API的容器执行器"""
    
    def __init__(self, config: ExecutorConfig, k8s_config: Dict[str, Any]):
        super().__init__(config)
        self.k8s_config = k8s_config
        self.namespace = k8s_config.get("namespace", "code-sandbox")
        self.api_server = k8s_config.get("api_server")
        self.token = k8s_config.get("token")
        self.ca_cert = k8s_config.get("ca_cert")
        
        # 容器镜像映射
        self.container_images = {
            CodeType.PYTHON: "python:3.9-alpine",
            CodeType.JAVASCRIPT: "node:16-alpine",
            CodeType.BASH: "alpine:latest",
            CodeType.JAVA: "openjdk:11-alpine"
        }
        
        # 执行命令映射
        self.exec_commands = {
            CodeType.PYTHON: ["python", "/workspace/code.py"],
            CodeType.JAVASCRIPT: ["node", "/workspace/code.js"],
            CodeType.BASH: ["sh", "/workspace/code.sh"],
            CodeType.JAVA: ["sh", "-c", "cd /workspace && javac code.java && java Main"]
        }
        
        self.active_pods: Dict[str, str] = {}  # task_id -> pod_name
    
    async def prepare_environment(self, request: CodeExecutionRequest) -> Dict[str, Any]:
        """准备容器执行环境"""
        # 生成唯一的pod名称
        pod_name = f"code-exec-{datetime.now(timezone.utc).strftime('%Y%m%d%H%M%S')}-{id(request) % 10000}"
        
        # 获取文件扩展名
        file_extensions = {
            CodeType.PYTHON: "code.py",
            CodeType.JAVASCRIPT: "code.js", 
            CodeType.BASH: "code.sh",
            CodeType.JAVA: "code.java"
        }
        
        code_filename = file_extensions.get(request.code_type, "code.txt")
        
        # 将代码编码为base64
        code_base64 = base64.b64encode(request.code.encode('utf-8')).decode('ascii')
        
        return {
            "pod_name": pod_name,
            "code_filename": code_filename,
            "code_base64": code_base64,
            "image": self.container_images.get(request.code_type),
            "command": self.exec_commands.get(request.code_type)
        }
    
    async def apply_security_constraints(self, request: CodeExecutionRequest, env_info: Dict[str, Any]) -> Dict[str, Any]:
        """应用安全约束"""
        return {
            "security_context": {
                "runAsNonRoot": True,
                "runAsUser": 1000,
                "readOnlyRootFilesystem": True,
                "allowPrivilegeEscalation": False,
                "capabilities": {
                    "drop": ["ALL"]
                }
            },
            "resource_limits": {
                "memory": f"{request.memory_limit_mb}Mi",
                "cpu": f"{request.cpu_limit}",
                "ephemeral-storage": "100Mi"
            },
            "resource_requests": {
                "memory": f"{max(16, request.memory_limit_mb // 2)}Mi", 
                "cpu": f"{max(0.1, request.cpu_limit / 2)}",
                "ephemeral-storage": "50Mi"
            },
            "network_policy": "restricted",
            "volume_mounts": [
                {
                    "name": "workspace",
                    "mountPath": "/workspace"
                },
                {
                    "name": "tmp",
                    "mountPath": "/tmp"
                }
            ]
        }
    
    def _create_pod_manifest(self, request: CodeExecutionRequest, env_info: Dict[str, Any], security_config: Dict[str, Any]) -> Dict[str, Any]:
        """创建Pod配置清单"""
        pod_name = env_info["pod_name"]
        
        # 创建ConfigMap引用来存储代码
        configmap_name = f"{pod_name}-code"
        
        pod_manifest = {
            "apiVersion": "v1",
            "kind": "Pod",
            "metadata": {
                "name": pod_name,
                "namespace": self.namespace,
                "labels": {
                    "app": "code-sandbox",
                    "executor": "container",
                    "code-type": request.code_type.value,
                    "user": request.user_info.user_id
                },
                "annotations": {
                    "security-level": request.security_level.value
                }
            },
            "spec": {
                "restartPolicy": "Never",
                "activeDeadlineSeconds": request.timeout_seconds,
                "securityContext": security_config["security_context"],
                "containers": [
                    {
                        "name": "code-runner",
                        "image": env_info["image"],
                        "command": env_info["command"],
                        "resources": {
                            "limits": security_config["resource_limits"],
                            "requests": security_config["resource_requests"]
                        },
                        "securityContext": {
                            "runAsNonRoot": True,
                            "runAsUser": 1000,
                            "readOnlyRootFilesystem": True,
                            "allowPrivilegeEscalation": False
                        },
                        "volumeMounts": [
                            {
                                "name": "code-volume",
                                "mountPath": "/workspace",
                                "readOnly": True
                            },
                            {
                                "name": "tmp-volume",
                                "mountPath": "/tmp"
                            }
                        ],
                        "env": [
                            {"name": key, "value": value} 
                            for key, value in request.environment_vars.items()
                        ]
                    }
                ],
                "volumes": [
                    {
                        "name": "code-volume",
                        "configMap": {
                            "name": configmap_name,
                            "defaultMode": 0o755
                        }
                    },
                    {
                        "name": "tmp-volume",
                        "emptyDir": {
                            "sizeLimit": "50Mi"
                        }
                    }
                ]
            }
        }
        
        return pod_manifest
    
    def _create_configmap_manifest(self, env_info: Dict[str, Any]) -> Dict[str, Any]:
        """创建ConfigMap配置清单"""
        pod_name = env_info["pod_name"]
        configmap_name = f"{pod_name}-code"
        
        return {
            "apiVersion": "v1",
            "kind": "ConfigMap",
            "metadata": {
                "name": configmap_name,
                "namespace": self.namespace
            },
            "binaryData": {
                env_info["code_filename"]: env_info["code_base64"]
            }
        }
    
    async def _k8s_api_request(self, method: str, path: str, data: Optional[Dict] = None) -> Dict[str, Any]:
        """发送Kubernetes API请求"""
        url = f"{self.api_server}{path}"
        headers = {
            "Authorization": f"Bearer {self.token}",
            "Content-Type": "application/json"
        }
        
        async with aiohttp.ClientSession() as session:
            kwargs = {
                "headers": headers,
                "ssl": False  # 在生产环境中应该验证SSL
            }
            
            if data:
                kwargs["json"] = data
            
            async with session.request(method, url, **kwargs) as response:
                response_text = await response.text()
                
                if response.status >= 400:
                    logger.error(f"Kubernetes API错误: {response.status} - {response_text}")
                    raise Exception(f"Kubernetes API错误: {response.status}")
                
                return json.loads(response_text) if response_text else {}
    
    async def execute_code(self, request: CodeExecutionRequest, env_info: Dict[str, Any], security_config: Dict[str, Any]) -> ExecutionResult:
        """执行代码"""
        pod_name = env_info["pod_name"]
        
        try:
            # 1. 创建ConfigMap
            configmap_manifest = self._create_configmap_manifest(env_info)
            await self._k8s_api_request(
                "POST",
                f"/api/v1/namespaces/{self.namespace}/configmaps",
                configmap_manifest
            )
            logger.debug(f"已创建ConfigMap: {configmap_manifest['metadata']['name']}")
            
            # 2. 创建Pod
            pod_manifest = self._create_pod_manifest(request, env_info, security_config)
            await self._k8s_api_request(
                "POST", 
                f"/api/v1/namespaces/{self.namespace}/pods",
                pod_manifest
            )
            logger.debug(f"已创建Pod: {pod_name}")
            
            # 记录活跃的Pod
            self.active_pods[str(id(request))] = pod_name
            
            # 3. 等待Pod完成
            result = await self._wait_for_pod_completion(pod_name, request.timeout_seconds)
            
            # 4. 获取日志
            if result.status in [TaskStatus.COMPLETED, TaskStatus.FAILED]:
                try:
                    logs = await self._get_pod_logs(pod_name)
                    result.output = logs
                except Exception as e:
                    logger.warning(f"获取Pod日志失败: {e}")
            
            return result
            
        except Exception as e:
            logger.error(f"容器执行失败: {e}")
            return ExecutionResult(
                task_id="",
                status=TaskStatus.FAILED,
                error=f"容器执行失败: {str(e)}",
                output=None,
                return_code=1,
                execution_time=0.0,
                memory_usage=0,
                cpu_usage=0.0,
                start_time=None,
                end_time=None
            )
    
    async def _wait_for_pod_completion(self, pod_name: str, timeout_seconds: int) -> ExecutionResult:
        """等待Pod执行完成"""
        start_time = datetime.now(timezone.utc)
        timeout_time = start_time + timedelta(seconds=timeout_seconds)
        
        while datetime.now(timezone.utc) < timeout_time:
            try:
                # 获取Pod状态
                pod_info = await self._k8s_api_request(
                    "GET",
                    f"/api/v1/namespaces/{self.namespace}/pods/{pod_name}"
                )
                
                status = pod_info.get("status", {})
                phase = status.get("phase")
                
                if phase == "Succeeded":
                    return ExecutionResult(
                        task_id="",
                        status=TaskStatus.COMPLETED,
                        return_code=0,
                        output=None,
                        error="",
                        execution_time=0.0,
                        memory_usage=0,
                        cpu_usage=0.0,
                        start_time=None,
                        end_time=None
                    )
                elif phase == "Failed":
                    container_statuses = status.get("containerStatuses", [])
                    error_msg = "容器执行失败"
                    
                    if container_statuses:
                        terminated = container_statuses[0].get("state", {}).get("terminated", {})
                        if terminated:
                            error_msg = terminated.get("reason", error_msg)
                            return_code = terminated.get("exitCode", 1)
                        else:
                            return_code = 1
                    else:
                        return_code = 1
                    
                    return ExecutionResult(
                        task_id="",
                        status=TaskStatus.FAILED,
                        error=error_msg,
                        return_code=return_code,
                        output=None,
                        execution_time=0.0,
                        memory_usage=0,
                        cpu_usage=0.0,
                        start_time=None,
                        end_time=None
                    )
                elif phase in ["Pending", "Running"]:
                    # 继续等待
                    await asyncio.sleep(1)
                else:
                    return ExecutionResult(
                        task_id="",
                        status=TaskStatus.FAILED,
                        error=f"未知的Pod状态: {phase}",
                        return_code=1,
                        output=None,
                        execution_time=0.0,
                        memory_usage=0,
                        cpu_usage=0.0,
                        start_time=None,
                        end_time=None
                    )
                    
            except Exception as e:
                logger.error(f"检查Pod状态失败: {e}")
                await asyncio.sleep(1)
        
        # 超时
        return ExecutionResult(
            task_id="",
            status=TaskStatus.TIMEOUT,
            error="Pod执行超时",
            return_code=-1,
            output=None,
            execution_time=0.0,
            memory_usage=0,
            cpu_usage=0.0,
            start_time=None,
            end_time=None
        )
    
    async def _get_pod_logs(self, pod_name: str) -> str:
        """获取Pod日志"""
        try:
            logs_response = await self._k8s_api_request(
                "GET",
                f"/api/v1/namespaces/{self.namespace}/pods/{pod_name}/log"
            )
            return logs_response if isinstance(logs_response, str) else ""
        except Exception as e:
            logger.error(f"获取Pod日志失败: {e}")
            return ""
    
    async def cleanup_environment(self, env_info: Dict[str, Any]) -> None:
        """清理执行环境"""
        pod_name = env_info.get("pod_name")
        configmap_name = f"{pod_name}-code"
        
        if pod_name:
            try:
                # 删除Pod
                await self._k8s_api_request(
                    "DELETE",
                    f"/api/v1/namespaces/{self.namespace}/pods/{pod_name}"
                )
                logger.debug(f"已删除Pod: {pod_name}")
            except Exception as e:
                logger.error(f"删除Pod失败: {e}")
            
            try:
                # 删除ConfigMap
                await self._k8s_api_request(
                    "DELETE",
                    f"/api/v1/namespaces/{self.namespace}/configmaps/{configmap_name}"
                )
                logger.debug(f"已删除ConfigMap: {configmap_name}")
            except Exception as e:
                logger.error(f"删除ConfigMap失败: {e}")
        
        # 从活跃Pod列表中移除
        for key, value in list(self.active_pods.items()):
            if value == pod_name:
                del self.active_pods[key]
                break 